#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "byte_code.h"
#include "obj.h"
#include "compiler.h"

static void expression();
static void grouping();
static void number();
static void unary();
static void binary();
static void advance();
static void consume(TokenType type, const char* msg);
static void parsePrece(Precedence prece);
static ParseRule getRule(TokenType type);
static void emitConstant(val);
static bool match(TokenType type);
static void declaration();
static void varDeclare();
static void statement();
static void printStmt();
static void expressionStmt();
static void variable();
static void string();
static void block();
static int localPosGet(Token token);
static void scopeEnd();
static void ifStmt();
static void whileStmt();
static void and_();
static void or_();
static void emitJmp(OpCode code);
static void dispatchJmp(int offset);
static void globalDeclare(Value val);
static void funDeclare();
static void call();
static void returnStmt();

extern const char* opCodeStrs;
Parser parser;
Chunk* chunk;
ObjFunction* globalFun;
int ebp; // 栈顶在locals中的索引
HashTable strings;
ParseRule rules[] = {
    [TOKEN_OR] = {NULL, or_, PRECE_OR},
    [TOKEN_AND] = {NULL, and_, PRECE_AND},
    [TOKEN_LESS] = {NULL, binary, PRECE_COMPARE},
    [TOKEN_LESS_EQUAL] = {NULL, binary, PRECE_COMPARE},
    [TOKEN_GREATER] = {NULL, binary, PRECE_COMPARE},
    [TOKEN_GREATER_EQUAL] = {NULL, binary, PRECE_COMPARE},
    [TOKEN_EQUAL_EQUAL] = {NULL, binary, PRECE_COMPARE},
    [TOKEN_BANG_EQUAL] = {NULL, binary, PRECE_COMPARE},
    [TOKEN_PLUS] = {NULL, binary, PRECE_TERM},
    [TOKEN_MINUS] = {unary, binary, PRECE_TERM},
    [TOKEN_STAR] = {NULL, binary, PRECE_FACTOR},
    [TOKEN_SLASH] = {NULL, binary, PRECE_FACTOR},
    [TOKEN_LEFT_PAREN] = {grouping, call, PRECE_CALL},
    [TOKEN_NUMBER] = {number, NULL, PRECE_NONE},
    [TOKEN_IDENTIFIER] = {variable, NULL, PRECE_NONE},
    [TOKEN_STRING] = {string, NULL, PRECE_NONE},
    [TOKEN_BANG] = {unary, NULL, PRECE_NONE},
    [TOKEN_EOF] = {NULL, NULL, PRECE_NONE},
};

static int localCompare(Local* l1, Local* l2) {
    Token t1 = l1->token;
    Token t2 = l2->token;
    if(t1.len != t2.len || t1.type != t2.type) return false;
    return memcmp(t1.start, t2.start, t1.len);
}

void parserInit(Parser* parser) {
    vectorInitCompare(&parser->locals, sizeof(Local), localCompare);
    parser->scopeDepth = 0;
}

void compile(const char* src) {
    parserInit(&parser);
    initScanner(src);
    hashTableInit(&strings, objStrHash, objStrCompare);
    globalFun = objFunctionCreate("", 0, &strings);
    chunk = &globalFun->chunk;
    ebp = 0;

    advance();
    while (parser.current.type != TOKEN_EOF)
    {
        declaration();
    }
    // printf("end:\n");
    // printToken(parser.current);
    consume(TOKEN_EOF, "token未解析完！");
    chunkAddCode(chunk, OP_RETURN, parser.current.line);
}

static void declaration() {
    if(match(TOKEN_VAR)) {
        varDeclare();
    } else if(match(TOKEN_FUN)) {
        funDeclare();
    } else {
        statement();
    }
}

static void funDeclare() {
    Token identifier = parser.current;
    ObjFunction* fun = objFunctionCreate(identifier.start, identifier.len, &strings);
    chunk = &fun->chunk;
    ebp = parser.locals.size;
    advance();
    consume(TOKEN_LEFT_PAREN, "函数声明需要(");
    parser.scopeDepth++;
    int arity = 0;
    while(!match(TOKEN_RIGHT_PAREN)) {
        Local local;
        local.token = parser.current;
        local.scopeDepth = parser.scopeDepth;
        vectorAppend(&parser.locals, &local);
        arity++;
        advance();
        if(!match(TOKEN_COMMA)) {
            consume(TOKEN_RIGHT_PAREN, "函数声明需要)");
            break;
        }
    }
    fun->arity = arity;
    // consume(TOKEN_RIGHT_PAREN, "函数声明需要)");
    consume(TOKEN_LEFT_BRACE, "函数声明需要{");
    
    block();
    chunkAddCode(chunk, OP_NIL, parser.current.line);
    chunkAddCode(chunk, OP_RETURN, parser.current.line);
    parser.scopeDepth--;
    chunk = &globalFun->chunk;
    // push fun instruction
    Value funVal;
    funVal.type = VAL_FUNCTION;
    funVal.val.obj = fun;
    emitConstant(funVal, identifier.line);
    Value varName;
    varName.type = VAL_STRING;
    varName.val.obj = fun->name;
    globalDeclare(varName);
}

static void call() {
    int args = 0;
    while(!match(TOKEN_RIGHT_PAREN)) {
        expression();
        args++;
        if(!match(TOKEN_COMMA)) {
            consume(TOKEN_RIGHT_PAREN, "调用表达式需要)");
            break;
        }
    }
    chunkAddCode(chunk, OP_CALL, parser.current.line);
    chunkAddCode(chunk, args, parser.current.line);
}

static void returnStmt() {
    if(!match(TOKEN_SEMICOLON)) {
        expression();
        consume(TOKEN_SEMICOLON, "返回语句需要;");
    } else {
        chunkAddCode(chunk, OP_NIL, parser.current.line);
    }
    chunkAddCode(chunk, OP_RETURN, parser.current.line);
}

static void varDeclare() {
    Token identifier = parser.current;
    advance();
    if(match(TOKEN_EQUAL)) {
        expression();
    } else {
        chunkAddCode(chunk, OP_NIL, identifier.line);
    }
    // printf("varDeclare %.*s\n", parser.current.len, parser.current.start);
    consume(TOKEN_SEMICOLON, "变量声明应该以分号结束！");
    if(parser.scopeDepth > 0) {
        Local local;
        local.token = identifier;
        local.scopeDepth = parser.scopeDepth;
        vectorAppend(&parser.locals, &local);
        return;
    }
    ObjString* varNameStr = objStringCreate(identifier.start, identifier.len, &strings);
    Value varName;
    varName.type = VAL_STRING;
    varName.val.obj = varNameStr;
    globalDeclare(varName);
}

static void statement() {
    if(match(TOKEN_PRINT)) {
        printStmt();
    } else if(match(TOKEN_LEFT_BRACE)) {
        parser.scopeDepth++;
        block();
        scopeEnd();
    } else if(match(TOKEN_IF)) {
        ifStmt();
    } else if(match(TOKEN_WHILE)) {
        whileStmt();
    } else if(match(TOKEN_RETURN)) {
        returnStmt();
    } else {
        expressionStmt();
    }
}

static void block() {
    // printf("block\n");
    
    while (parser.current.type != TOKEN_EOF && parser.current.type != TOKEN_RIGHT_BRACE)
    {
        declaration();
    }
    consume(TOKEN_RIGHT_BRACE, "块表达式应该以}结束");
    
}

static void whileStmt() {
    consume(TOKEN_LEFT_PAREN, "expect ( after while");
    int loopPos = chunk->codes.size;
    expression();
    consume(TOKEN_RIGHT_PAREN, "expect ) after while expression");
    emitJmp(OP_JMP_IF_FALSE);
    int jmpPos = chunk->codes.size - 1;
    // then start
    chunkAddCode(chunk, OP_POP, parser.current.line);
    statement();
    chunkAddCode(chunk, OP_LOOP, parser.current.line);
    // ip = chunk->codes.size + 1; ip - ? = loopPos - 1
    int steps = chunk->codes.size - loopPos + 2;
    chunkAddCode(chunk, (steps >> 8) & 0xff, parser.current.line);
    chunkAddCode(chunk, steps & 0xff, parser.current.line);
    // then end
    dispatchJmp(jmpPos);
    // else start
    chunkAddCode(chunk, OP_POP, parser.current.line);
}

static void ifStmt() {
    consume(TOKEN_LEFT_PAREN, "expect ( after if");
    expression();
    consume(TOKEN_RIGHT_PAREN, "expect ) after if expression");
    emitJmp(OP_JMP_IF_FALSE);
    int thenPos = chunk->codes.size - 1;
    // then start
    chunkAddCode(chunk, OP_POP, parser.current.line);
    statement();
    emitJmp(OP_JMP);
    int elsePos = chunk->codes.size - 1;
    // else start
    dispatchJmp(thenPos);
    chunkAddCode(chunk, OP_POP, parser.current.line);
    if(match(TOKEN_ELSE)) {
        statement();
    }
    // else end
    dispatchJmp(elsePos);
}

static void printStmt() {
    expression();
    consume(TOKEN_SEMICOLON, "print语句以分号结束");
    chunkAddCode(chunk, OP_PRINT, parser.current.line);
}

static void expressionStmt() {
    expression();
    consume(TOKEN_SEMICOLON, "表达式语句以分号结束");
    chunkAddCode(chunk, OP_POP, parser.current.line);
}

static void expression() {
    parsePrece(PRECE_ASSIGNMENT);
}

static void grouping() {
    expression();
    consume(TOKEN_RIGHT_PAREN, "()表达式应该以)结尾！");
}

static void variable() {
    Token identifier = parser.previous;
    OpCode opCode = OP_VARIABLE_GET;
    if(match(TOKEN_EQUAL)) {
        expression();
        opCode = OP_ASSIGNMENT;
    }
    if(parser.scopeDepth > 0) {
        int idx = localPosGet(identifier);
        if(idx != -1) {
            OpCode localCode = opCode == OP_VARIABLE_GET ? OP_LOCAL_GET : OP_LOCAL_ASSIGN;
            chunkAddCode(chunk, localCode, identifier.line);
            chunkAddCode(chunk, idx, identifier.line);
            return;
        }
    }
    ObjString* varNameStr = objStringCreate(identifier.start, identifier.len, &strings);
    Value varName;
    varName.type = VAL_STRING;
    varName.val.obj = varNameStr;
    int idx = chunkAddConstant(chunk, varName);
    chunkAddCode(chunk, opCode, identifier.line);
    chunkAddCode(chunk, idx, identifier.line);
}

static void string() {
    // printf("string\n");
    // printf("%.*s\n", parser.previous.len - 2, parser.previous.start + 1);
    ObjString* str = objStringCreate(parser.previous.start + 1, 
        parser.previous.len - 2, &strings);
    // printf("string %p, %s\n", str, str->val);
    Value val;
    val.type = VAL_STRING;
    val.val.obj = str;
    emitConstant(val, parser.previous.line);
    // printf("string end\n");
}

static void number() {
    double num = strtod(parser.previous.start, NULL);
    // printf("number: %.2f\n", num);
    Value val = chunkNum2Val(num);
    emitConstant(val, parser.previous.line);
}

static void unary() {
    Token token = parser.previous;
    parsePrece(PRECE_UNARY);
    switch (token.type)
    {
    case TOKEN_MINUS:
        chunkAddCode(chunk, OP_NEGATE, token.line);
        break;
    case TOKEN_BANG:
        chunkAddCode(chunk, OP_NOT, token.line);
        break;
    
    default:
        break;
    }
}

static void binary() {
    Token token = parser.previous;
    ParseRule rule = getRule(token.type);
    parsePrece(rule.prece + 1);
    switch (token.type)
    {
    case TOKEN_PLUS:
        chunkAddCode(chunk, OP_ADD, token.line);
        break;
    case TOKEN_MINUS:
        chunkAddCode(chunk, OP_SUBSTRACT, token.line);
        break;
    case TOKEN_STAR:
        chunkAddCode(chunk, OP_MULTIPLY, token.line);
        break;
    case TOKEN_SLASH:
        chunkAddCode(chunk, OP_DIVIDE, token.line);
        break;
    case TOKEN_LESS:
        chunkAddCode(chunk, OP_LESS, token.line);
        break;
    case TOKEN_GREATER:
        chunkAddCode(chunk, OP_GREATER, token.line);
        break;
    case TOKEN_LESS_EQUAL:
        chunkAddCode(chunk, OP_GREATER, token.line);
        chunkAddCode(chunk, OP_NOT, token.line);
        break;
    case TOKEN_GREATER_EQUAL:
        chunkAddCode(chunk, OP_LESS, token.line);
        chunkAddCode(chunk, OP_NOT, token.line);
        break;
    case TOKEN_EQUAL:
        chunkAddCode(chunk, OP_EQUAL, token.line);
        break;
    case TOKEN_BANG_EQUAL:
        chunkAddCode(chunk, OP_EQUAL, token.line);
        chunkAddCode(chunk, OP_NOT, token.line);
        break;
    
    default:
        break;
    }
}

static void and_() {
    emitJmp(OP_JMP_IF_FALSE);
    int jmpPos = chunk->codes.size - 1;
    // then
    chunkAddCode(chunk, OP_POP, parser.current.line);
    parsePrece(PRECE_AND + 1);
    // else
    dispatchJmp(jmpPos);
}

static void or_() {
    emitJmp(OP_JMP_IF_TRUE);
    int jmpPos = chunk->codes.size - 1;
    // then
    chunkAddCode(chunk, OP_POP, parser.current.line);
    parsePrece(PRECE_AND);
    // else
    dispatchJmp(jmpPos);
}

static void advance() {
    parser.previous = parser.current;
    for(;;) {
        Token token = scanToken();
        // printToken(token);
        if(token.type == TOKEN_ERROR) {
            fprintf(stderr, "error: %s\n", token.start);
        } else {
            parser.current = token;
            break;
        }
    }
}

static void consume(TokenType type, const char* msg) {
    if(parser.current.type != type) {
        fprintf(stderr, "line[%d]: %s\n", parser.current.line, msg);
    } else {
        advance();
    }
}

static void parsePrece(Precedence prece) {
    advance();
    ParseRule rulePre = getRule(parser.previous.type);
    if(rulePre.prefix) {
        rulePre.prefix();
    } else {
        fprintf(stderr, "tokenType: %d, 缺少解析方法！", parser.previous.type);
        return;
    }
    while(prece <= getRule(parser.current.type).prece) {
        advance();
        getRule(parser.previous.type).infix();
    }
}

static ParseRule getRule(TokenType type) {
    return rules[type];
}

static void emitConstant(Value val, int line) {
    chunkAddCode(chunk, OP_CONSTANT, line);
    int idx = chunkAddConstant(chunk, val);
    chunkAddCode(chunk, idx, line);
}

static bool match(TokenType type) {
    if(parser.current.type == type) {
        advance();
        return true;
    } else {
        return false;
    }
}

static int localPosGet(Token token) {
    Local local = {token};
    return vectorFindLast(&parser.locals, &local) - ebp;
}

static void scopeEnd() {
    while (parser.locals.size > 0 && 
        ((Local*)vectorGet(&parser.locals, parser.locals.size -1))->scopeDepth >= parser.scopeDepth)
    {
        vectorRemove(&parser.locals, parser.locals.size - 1);
        chunkAddCode(chunk, OP_POP, parser.current.line);
    }
    parser.scopeDepth--;
}

static void emitJmp(OpCode code) {
    chunkAddCode(chunk, code, parser.current.line);
    chunkAddCode(chunk, 0, parser.current.line);
    chunkAddCode(chunk, 0, parser.current.line);
}

// offset表示jmp类3条指令的最后一条的索引
static void dispatchJmp(int offset) {
    // -1 是由于虚拟机循环体最后都要vm->ip++
    int increment = chunk->codes.size - offset - 1;
    // &的优先级高于>>
    chunkUpdateCode(chunk, (increment >> 8) & 0xff, offset - 1);
    chunkUpdateCode(chunk, increment & 0xff, offset);
}

static void globalDeclare(Value val) {
    int idx = chunkAddConstant(chunk, val);
    chunkAddCode(chunk, OP_VARIABLE_DECLARE, parser.current.line);
    chunkAddCode(chunk, idx, parser.current.line);
    chunkAddCode(chunk, OP_POP, parser.current.line);
}
